package com.watson.guava.config;

import com.google.common.util.concurrent.RateLimiter;
import com.watson.guava.annotation.RequestLimiter;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Component;
import org.springframework.web.filter.GenericFilterBean;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;
import org.springframework.web.servlet.ModelAndView;
import org.springframework.web.servlet.config.annotation.InterceptorRegistry;
import org.springframework.web.servlet.handler.HandlerInterceptorAdapter;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.TimeUnit;

/**
 * // 拦截器实现限流
 *
 * @author: llun
 * @date: 2022/2/15 15:59
 * @version: 1.0
 */
@Component
public class RequestLimiterInterceptor extends HandlerInterceptorAdapter {

    public enum LimitType {
        DROP,	//丢弃
        WAIT	//等待
    }

    /**
     * Guava 开源工具限流工具类
     * 限流器
     */
    private RateLimiter limiter;

    /**
     * 限流方式
     */
    private LimitType limitType = LimitType.DROP;

    public RequestLimiterInterceptor() {
        this.limiter = RateLimiter.create(10);
    }

    /**
     * @param tps	限流（每秒处理量）
     * @param limitType
     */
    public RequestLimiterInterceptor(int tps, RequestLimiterInterceptor.LimitType limitType) {
        this.limiter = RateLimiter.create(tps);
        this.limitType = limitType;
    }

    /**
     * @param permitsPerSecond	每秒新增的令牌数
     * @param limitType	限流类型
     */
    public RequestLimiterInterceptor(double permitsPerSecond, RequestLimiterInterceptor.LimitType limitType) {
        this.limiter = RateLimiter.create(permitsPerSecond, 1000, TimeUnit.MILLISECONDS);
        this.limitType = limitType;
    }

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (limitType.equals(LimitType.DROP)) {
            //尝试获取一个令牌,立即返回
            if (limiter.tryAcquire()) {

                return super.preHandle(request, response, handler);
            }
          return  false;
        }
	    else {
	    	//获取令牌 如果没有令牌则一直等待,返回等待的时间(单位为秒),没有被限流则直接返回0.0:
			double count = limiter.acquire();
	    	return super.preHandle(request, response, handler);
	    }
//        throw new RuntimeException("服务器达到请求达到上限，限流生效");//达到限流后，往页面提示的错误信息。
    }

    @Override
    public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView) throws Exception {
        super.postHandle(request, response, handler, modelAndView);
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
        super.afterCompletion(request, response, handler, ex);
    }


}
